
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/err.h>


#define CHECK_AND_RETURN(condition, ret_val, message)
    do {
        if (!(condition)) {
            fprintf(stderr, "Error: %s\n", message);
            ERR_print_errors_fp(stderr);
            return ret_val;
        }
    } while(0)


void print_hex(const char *label, const uint8_t *data, uint32_t len) {
    printf("%s: ", label);
    for (uint32_t i = 0; i < len; i++) {
        printf("%02x", data[i]);
    }
    printf("\n");
}


int aes_cbc_encrypt(const uint8_t *plaintext, uint32_t plaintext_len,
                    const uint8_t *key, uint32_t key_len,
                    const uint8_t *iv, uint32_t iv_len,
                    uint8_t *ciphertext, uint32_t *ciphertext_len) {
    EVP_CIPHER_CTX *ctx = NULL;
    int len = 0;
    int total_len = 0;
    const EVP_CIPHER *cipher = NULL;
    switch (key_len) {
        case 16: cipher = EVP_aes_128_cbc(); break;
        case 24: cipher = EVP_aes_192_cbc(); break;
        case 32: cipher = EVP_aes_256_cbc(); break;
        default:
            fprintf(stderr, "Error: Invalid key length %u\n", key_len);
            return 0;
    }
    
    ctx = EVP_CIPHER_CTX_new();
    CHECK_AND_RETURN(ctx != NULL, 0, "Failed to create cipher context");
    
    CHECK_AND_RETURN(EVP_EncryptInit_ex(ctx, cipher, NULL, key, iv) == 1, 
                     0, "Failed to initialize encryption");
    
    CHECK_AND_RETURN(EVP_CIPHER_CTX_set_padding(ctx, 1) == 1,
                     0, "Failed to set padding");

    CHECK_AND_RETURN(EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len) == 1, 
                     0, "Failed to encrypt data");
    total_len = len;

    CHECK_AND_RETURN(EVP_EncryptFinal_ex(ctx, ciphertext + len, &len) == 1, 
                     0, "Failed to finalize encryption");
    total_len += len;
    
    *ciphertext_len = total_len;

    EVP_CIPHER_CTX_free(ctx);
    
    return 1;
}

int aes_cbc_decrypt(const uint8_t *ciphertext, uint32_t ciphertext_len,
                    const uint8_t *key, uint32_t key_len,
                    const uint8_t *iv, uint32_t iv_len,
                    uint8_t *plaintext, uint32_t *plaintext_len) {
    EVP_CIPHER_CTX *ctx = NULL;
    int len = 0;
    int total_len = 0;
    const EVP_CIPHER *cipher = NULL;

    switch (key_len) {
        case 16: cipher = EVP_aes_128_cbc(); break;
        case 24: cipher = EVP_aes_192_cbc(); break;
        case 32: cipher = EVP_aes_256_cbc(); break;
        default:
            fprintf(stderr, "Error: Invalid key length %u\n", key_len);
            return 0;
    }

    ctx = EVP_CIPHER_CTX_new();
    CHECK_AND_RETURN(ctx != NULL, 0, "Failed to create cipher context");

    CHECK_AND_RETURN(EVP_DecryptInit_ex(ctx, cipher, NULL, key, iv) == 1, 
                     0, "Failed to initialize decryption");
    
    CHECK_AND_RETURN(EVP_CIPHER_CTX_set_padding(ctx, 1) == 1,
                     0, "Failed to set padding");
    
    CHECK_AND_RETURN(EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len) == 1, 
                     0, "Failed to decrypt data");
    total_len = len;

    CHECK_AND_RETURN(EVP_DecryptFinal_ex(ctx, plaintext + len, &len) == 1, 
                     0, "Failed to finalize decryption");
    total_len += len;
    
    *plaintext_len = total_len;

    EVP_CIPHER_CTX_free(ctx);
    
    return 1;
}

int main(void)
{

    OpenSSL_add_all_algorithms();
    ERR_load_crypto_strings();
    
    uint8_t data[10] = {0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14, 0x1c, 0x14};
    uint8_t iv[16] = {0};
    uint8_t key[16] = {0};
    uint32_t dataLen = sizeof(data);
    uint8_t cipherText[100];
    uint8_t plainText[100];
    uint32_t cipherTextLen = 0;
    uint32_t plainTextLen = 0;
    int ret = 1;
    
    print_hex("Plain text to be encrypted", data, dataLen);

    if (!aes_cbc_encrypt(data, dataLen, key, sizeof(key), iv, sizeof(iv), 
                         cipherText, &cipherTextLen)) {
        fprintf(stderr, "Encryption failed\n");
        goto CLEANUP;
    }
    
    print_hex("Cipher text", cipherText, cipherTextLen);

    if (!aes_cbc_decrypt(cipherText, cipherTextLen, key, sizeof(key), iv, sizeof(iv), 
                         plainText, &plainTextLen)) {
        fprintf(stderr, "Decryption failed\n");
        goto CLEANUP;
    }
    
    print_hex("Decrypted plain text", plainText, plainTextLen);

    if (plainTextLen != dataLen || memcmp(plainText, data, dataLen) != 0) {
        fprintf(stderr, "Plaintext comparison failed\n");
        goto CLEANUP;
    }
    
    printf("AES CBC test passed!\n");
    ret = 0; 

CLEANUP:
    EVP_cleanup();
    ERR_free_strings();
    
    return ret;
}